package org.loong.crypto.core.algorithm;

import java.security.NoSuchAlgorithmException;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.DERNull;
import org.bouncycastle.asn1.gm.GMObjectIdentifiers;
import org.bouncycastle.asn1.nist.NISTObjectIdentifiers;
import org.bouncycastle.asn1.oiw.OIWObjectIdentifiers;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.pkcs.RSASSAPSSparams;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x9.X9ObjectIdentifiers;
import org.loong.crypto.common.utils.ArrayUtils;
import org.loong.crypto.core.Algorithm;
import org.loong.crypto.core.AlgorithmFamily;

/**
 * The Signature algorithm.
 */
public final class SignatureAlgorithm extends Algorithm {

    /**
     * The digest algorithm of signature algorithm.
     */
    private MessageDigestAlgorithm digestAlgorithm;

    /**
     * Signature using SHA1WithDSA algorithm.
     */
    public static final SignatureAlgorithm SHA1_DSA = new SignatureAlgorithm("SHA1WithDSA", X9ObjectIdentifiers.id_dsa_with_sha1, MessageDigestAlgorithm.SHA1);

    /**
     * Signature using SHA224WithDSA algorithm.
     */
    public static final SignatureAlgorithm SHA224_DSA = new SignatureAlgorithm("SHA224WithDSA", NISTObjectIdentifiers.dsa_with_sha224, MessageDigestAlgorithm.SHA224);

    /**
     * Signature using SHA256WithDSA algorithm.
     */
    public static final SignatureAlgorithm SHA256_DSA = new SignatureAlgorithm("SHA256WithDSA", NISTObjectIdentifiers.dsa_with_sha256, MessageDigestAlgorithm.SHA256);

    /**
     * Signature using SHA384WithDSA algorithm.
     */
    public static final SignatureAlgorithm SHA384_DSA = new SignatureAlgorithm("SHA384WithDSA", NISTObjectIdentifiers.dsa_with_sha384, MessageDigestAlgorithm.SHA384);

    /**
     * Signature using SHA512WithDSA algorithm.
     */
    public static final SignatureAlgorithm SHA512_DSA = new SignatureAlgorithm("SHA512WithDSA", NISTObjectIdentifiers.dsa_with_sha512, MessageDigestAlgorithm.SHA512);

    /**
     * Signature using SHA1WithECDSA algorithm.
     */
    public static final SignatureAlgorithm SHA1_ECDSA = new SignatureAlgorithm("SHA1WithECDSA", X9ObjectIdentifiers.ecdsa_with_SHA1, MessageDigestAlgorithm.SHA1);

    /**
     * Signature using SHA224WithECDSA algorithm.
     */
    public static final SignatureAlgorithm SHA224_ECDSA = new SignatureAlgorithm("SHA224WithECDSA", X9ObjectIdentifiers.ecdsa_with_SHA224, MessageDigestAlgorithm.SHA224);

    /**
     * Signature using SHA256WithECDSA algorithm.
     */
    public static final SignatureAlgorithm SHA256_ECDSA = new SignatureAlgorithm("SHA256WithECDSA", X9ObjectIdentifiers.ecdsa_with_SHA256, MessageDigestAlgorithm.SHA256);

    /**
     * Signature using SHA384WithECDSA algorithm.
     */
    public static final SignatureAlgorithm SHA384_ECDSA = new SignatureAlgorithm("SHA384WithECDSA", X9ObjectIdentifiers.ecdsa_with_SHA384, MessageDigestAlgorithm.SHA384);

    /**
     * Signature using SHA512WithECDSA algorithm.
     */
    public static final SignatureAlgorithm SHA512_ECDSA = new SignatureAlgorithm("SHA512WithECDSA", X9ObjectIdentifiers.ecdsa_with_SHA512, MessageDigestAlgorithm.SHA512);

    /**
     * Signature using MD2WithRSA algorithm.
     */
    public static final SignatureAlgorithm MD2_RSA = new SignatureAlgorithm("MD2WithRSA", PKCSObjectIdentifiers.md2WithRSAEncryption, MessageDigestAlgorithm.MD2);

    /**
     * Signature using MD5WithRSA algorithm.
     */
    public static final SignatureAlgorithm MD5_RSA = new SignatureAlgorithm("MD5WithRSA", PKCSObjectIdentifiers.md5WithRSAEncryption, MessageDigestAlgorithm.MD5);

    /**
     * Signature using SHA1WithRSA algorithm.
     */
    public static final SignatureAlgorithm SHA1_RSA = new SignatureAlgorithm("SHA1WithRSA", PKCSObjectIdentifiers.sha1WithRSAEncryption, MessageDigestAlgorithm.SHA1);

    /**
     * Signature using SHA224WithRSA algorithm.
     */
    public static final SignatureAlgorithm SHA224_RSA = new SignatureAlgorithm("SHA224WithRSA", PKCSObjectIdentifiers.sha224WithRSAEncryption, MessageDigestAlgorithm.SHA224);

    /**
     * Signature using SHA256WithRSA algorithm.
     */
    public static final SignatureAlgorithm SHA256_RSA = new SignatureAlgorithm("SHA256WithRSA", PKCSObjectIdentifiers.sha256WithRSAEncryption, MessageDigestAlgorithm.SHA256);

    /**
     * Signature using SHA384WithRSA algorithm.
     */
    public static final SignatureAlgorithm SHA384_RSA = new SignatureAlgorithm("SHA384WithRSA", PKCSObjectIdentifiers.sha384WithRSAEncryption, MessageDigestAlgorithm.SHA384);

    /**
     * Signature using SHA512WithRSA algorithm.
     */
    public static final SignatureAlgorithm SHA512_RSA = new SignatureAlgorithm("SHA512WithRSA", PKCSObjectIdentifiers.sha512WithRSAEncryption, MessageDigestAlgorithm.SHA512);

    /**
     * Signature using SHA1withRSASSA-PSS algorithm.
     */
    public static final SignatureAlgorithm SHA1_RSASSA_PSS = new SignatureAlgorithm("SHA1WithRSASSA-PSS", PKCSObjectIdentifiers.id_RSASSA_PSS, MessageDigestAlgorithm.SHA1);

    /**
     * Signature using SHA224withRSASSA-PSS algorithm.
     */
    public static final SignatureAlgorithm SHA224_RSASSA_PSS =
            new SignatureAlgorithm("SHA224WithRSASSA-PSS", PKCSObjectIdentifiers.id_RSASSA_PSS, MessageDigestAlgorithm.SHA224);

    /**
     * Signature using SHA256withRSASSA-PSS algorithm.
     */
    public static final SignatureAlgorithm SHA256_RSASSA_PSS =
            new SignatureAlgorithm("SHA256WithRSASSA-PSS", PKCSObjectIdentifiers.id_RSASSA_PSS, MessageDigestAlgorithm.SHA256);

    /**
     * Signature using SHA384withRSASSA-PSS algorithm.
     */
    public static final SignatureAlgorithm SHA384_RSASSA_PSS =
            new SignatureAlgorithm("SHA384WithRSASSA-PSS", PKCSObjectIdentifiers.id_RSASSA_PSS, MessageDigestAlgorithm.SHA384);

    /**
     * Signature using SHA512withRSASSA-PSS algorithm.
     */
    public static final SignatureAlgorithm SHA512_RSASSA_PSS =
            new SignatureAlgorithm("SHA512WithRSASSA-PSS", PKCSObjectIdentifiers.id_RSASSA_PSS, MessageDigestAlgorithm.SHA512);

    /**
     * Signature using SHA1WithSM2 algorithm.
     */
    public static final SignatureAlgorithm SHA1_SM2 = new SignatureAlgorithm("SHA1WithSM2", GMObjectIdentifiers.sm2sign_with_sha1, MessageDigestAlgorithm.SHA1);

    /**
     * Signature using SHA256WithSM2 algorithm.
     */
    public static final SignatureAlgorithm SHA256_SM2 = new SignatureAlgorithm("SHA256WithSM2", GMObjectIdentifiers.sm2sign_with_sha256, MessageDigestAlgorithm.SHA256);

    /**
     * Signature using SM3WithSM2 algorithm.
     */
    public static final SignatureAlgorithm SM3_SM2 = new SignatureAlgorithm("SM3WithSM2", GMObjectIdentifiers.sm2sign_with_sm3, MessageDigestAlgorithm.SM3);

    /**
     * The supported signature algorithms.
     */
    public static final Map<String, SignatureAlgorithm> SUPPORTED_ALGORITHMS;

    private static final Map<String, RSASSAPSSparams> PARAMS;

    static {
        Map<String, SignatureAlgorithm> algs = new LinkedHashMap<>();
        // SHA with DSA
        algs.put(SHA1_DSA.getName(), SHA1_DSA);
        algs.put(SHA224_DSA.getName(), SHA224_DSA);
        algs.put(SHA256_DSA.getName(), SHA256_DSA);
        algs.put(SHA384_DSA.getName(), SHA384_DSA);
        algs.put(SHA512_DSA.getName(), SHA512_DSA);

        // SHA with ECDSA
        algs.put(SHA1_ECDSA.getName(), SHA1_ECDSA);
        algs.put(SHA224_ECDSA.getName(), SHA224_ECDSA);
        algs.put(SHA256_ECDSA.getName(), SHA256_ECDSA);
        algs.put(SHA384_ECDSA.getName(), SHA384_ECDSA);
        algs.put(SHA512_ECDSA.getName(), SHA512_ECDSA);

        // MD with RSA
        algs.put(MD2_RSA.getName(), MD2_RSA);
        algs.put(MD5_RSA.getName(), MD5_RSA);

        // SHA with RSA
        algs.put(SHA1_RSA.getName(), SHA1_RSA);
        algs.put(SHA224_RSA.getName(), SHA224_RSA);
        algs.put(SHA256_RSA.getName(), SHA256_RSA);
        algs.put(SHA384_RSA.getName(), SHA384_RSA);
        algs.put(SHA512_RSA.getName(), SHA512_RSA);

        // SHA with RSASSA-PSS
        algs.put(SHA1_RSASSA_PSS.getName(), SHA1_RSASSA_PSS);
        algs.put(SHA224_RSASSA_PSS.getName(), SHA224_RSASSA_PSS);
        algs.put(SHA256_RSASSA_PSS.getName(), SHA256_RSASSA_PSS);
        algs.put(SHA384_RSASSA_PSS.getName(), SHA384_RSASSA_PSS);
        algs.put(SHA512_RSASSA_PSS.getName(), SHA512_RSASSA_PSS);

        // SM2
        algs.put(SHA1_SM2.getName(), SHA1_SM2);
        algs.put(SHA256_SM2.getName(), SHA256_SM2);
        algs.put(SM3_SM2.getName(), SM3_SM2);
        SUPPORTED_ALGORITHMS = Collections.unmodifiableMap(algs);

        Map<String, RSASSAPSSparams> params = new ConcurrentHashMap<>();
        AlgorithmIdentifier sha1AlgId = new AlgorithmIdentifier(OIWObjectIdentifiers.idSHA1, DERNull.INSTANCE);
        params.put("SHA1WithRSASSA-PSS", createPSSParams(sha1AlgId, 20));

        AlgorithmIdentifier sha224AlgId = new AlgorithmIdentifier(NISTObjectIdentifiers.id_sha224, DERNull.INSTANCE);
        params.put("SHA224WithRSASSA-PSS", createPSSParams(sha224AlgId, 28));

        AlgorithmIdentifier sha256AlgId = new AlgorithmIdentifier(NISTObjectIdentifiers.id_sha256, DERNull.INSTANCE);
        params.put("SHA256WithRSASSA-PSS", createPSSParams(sha256AlgId, 32));

        AlgorithmIdentifier sha384AlgId = new AlgorithmIdentifier(NISTObjectIdentifiers.id_sha384, DERNull.INSTANCE);
        params.put("SHA384WithRSASSA-PSS", createPSSParams(sha384AlgId, 48));

        AlgorithmIdentifier sha512AlgId = new AlgorithmIdentifier(NISTObjectIdentifiers.id_sha512, DERNull.INSTANCE);
        params.put("SHA512WithRSASSA-PSS", createPSSParams(sha512AlgId, 64));
        PARAMS = Collections.unmodifiableMap(params);
    }

    /**
     * Signature algorithm family.
     */
    public static final class Family extends AlgorithmFamily<SignatureAlgorithm> {

        private static final long serialVersionUID = 1L;

        /**
         * RSA signature using MD2/4/5 hash.
         */
        public static final Family MD_RSA = new Family(MD2_RSA, MD5_RSA);

        /**
         * RSA signature using SHA1/224/256/384/512 hash.
         */
        public static final Family SHA_RSA =
                new Family(SHA1_RSA, SHA224_RSA, SHA256_RSA, SHA384_RSA, SHA512_RSA, SHA1_RSASSA_PSS, SHA224_RSASSA_PSS, SHA256_RSASSA_PSS, SHA384_RSASSA_PSS, SHA512_RSASSA_PSS);

        /**
         * RSA signature algorithms.
         */
        public static final Family RSA = new Family(ArrayUtils.concat(MD_RSA.toArray(new SignatureAlgorithm[] {}), SHA_RSA.toArray(new SignatureAlgorithm[] {})));

        /**
         * DSA signature using SHA1/224/256/384/512 hash.
         */
        public static final Family DSA = new Family(SHA1_DSA, SHA224_DSA, SHA256_DSA, SHA384_DSA, SHA512_DSA);

        /**
         * ECDSA signature using SHA1/224/256/384/512 hash.
         */
        public static final Family ECDSA = new Family(SHA1_ECDSA, SHA224_ECDSA, SHA256_ECDSA, SHA384_ECDSA, SHA512_ECDSA);

        /**
         * SM2 signature algorithms.
         */
        public static final Family SM2 = new Family(SHA1_SM2, SHA256_SM2, SM3_SM2);

        /**
         * Creates a new signature algorithm family.
         *
         * @param algs The signature algorithms of the family. Must not be {@code null}.
         */
        public Family(final SignatureAlgorithm... algs) {
            super(algs);
        }
    }

    /**
     * Creates a new signature algorithm name.
     * 
     * @param name the algorithm name. Must not be {@code null}.
     * @param oid the algorithm object identifier. Must not be {@code null}.
     * @param digestAlgorithm the digest algorithm of signature algorithm. Must not be {@code null}.
     */
    public SignatureAlgorithm(final String name, final ASN1ObjectIdentifier oid, final MessageDigestAlgorithm digestAlgorithm) {
        super(name, oid);

        this.digestAlgorithm = digestAlgorithm;
    }

    /**
     * Gets the value of digestAlgorithm.
     *
     * @return the value of digestAlgorithm
     */
    public final MessageDigestAlgorithm getDigestAlgorithm() {
        return digestAlgorithm;
    }

    /**
     * Gets signature algorithm identifier.
     *
     * @return the signature algorithm identifier
     */
    public AlgorithmIdentifier getSigAlgId() {
        final String algorithmName = this.getName();
        AlgorithmIdentifier sigAlgId = new AlgorithmIdentifier(getOID());
        if (PARAMS.containsKey(algorithmName)) {
            sigAlgId = new AlgorithmIdentifier(getOID(), (ASN1Encodable) PARAMS.get(algorithmName));
        }
        return sigAlgId;
    }

    /**
     * Create RSASSAPSSparams params.
     * 
     * @param hashAlgId the hashAlgId
     * @param saltSize the saltSize
     * @return the {@link RSASSAPSSparams}
     */
    private static RSASSAPSSparams createPSSParams(AlgorithmIdentifier hashAlgId, int saltSize) {
        return new RSASSAPSSparams(hashAlgId, new AlgorithmIdentifier(PKCSObjectIdentifiers.id_mgf1, hashAlgId), new ASN1Integer(saltSize), new ASN1Integer(1));
    }

    /**
     * Find a signature algorithm from the specified string.
     *
     * @param name the algorithm name. Must not be {@code null}.
     * @return the signature algorithm.
     * @throws NoSuchAlgorithmException the ${@link NoSuchAlgorithmException}
     */
    public static SignatureAlgorithm find(final String name) throws NoSuchAlgorithmException {
        SignatureAlgorithm algorithm = SUPPORTED_ALGORITHMS.get(name);
        if (algorithm == null) {
            throw new NoSuchAlgorithmException("Unknown signature algorithm type requested: " + name);
        }

        return algorithm;
    }

    /**
     * Find a signature algorithm from the specified string.
     *
     * @param algorithmId the signature algorithm id. Must not be {@code null}.
     * @return the signature algorithm.
     * @throws NoSuchAlgorithmException the ${@link NoSuchAlgorithmException}
     */
    public static SignatureAlgorithm findByAlgorithmId(final AlgorithmIdentifier algorithmId) throws NoSuchAlgorithmException {
        final ASN1ObjectIdentifier oid = algorithmId.getAlgorithm();
        return SUPPORTED_ALGORITHMS.values().stream().filter(algorithm -> Objects.equals(algorithm.getOID(), oid)).findFirst()
                .orElseThrow(() -> new NoSuchAlgorithmException("Unknown signature algorithm type requested: " + oid));
    }
}
